# -*- coding: utf-8 -*-
'''
Created on 13.05.2019

@author: yu03
'''
import numpy as np
import matplotlib.pyplot as plt
from scipy import signal
from scipy.optimize import curve_fit

P, Q = 1.36, 0.55 #Hanning Window
def gaus(x,a,x0,sigma):
    return a*np.exp(-(x-x0)**2/(2*sigma**2))

def FFT_interpolation_2(sig, tau0, zero_num, DC_num): ###zero_num=1e5, DC_num=3000
    zero_padd = np.zeros(int(zero_num))
    sig_padd = np.concatenate((sig, zero_padd))
    N_data = len(sig)
    freqline, sig_FFT, sig_magnitude, sig_phase = FFT_cal(sig_padd, tau0)
    freq_bin = 1/ tau0 / len(sig_padd)
#     print(freq_bin)
    m_k_num = sig_magnitude[DC_num:].argmax()+DC_num ### frequency bin 序数 (3000 for 1e5)
#     print(m_k_num)
    m_k = freqline[m_k_num] ### 换算为实际频率
#     print(m_k_num)
    X_m_k = sig_magnitude[m_k_num] ### magnitude (complex raw data)
    X_m_k_minus = sig_magnitude[m_k_num-1]
    X_m_k_plus = sig_magnitude[m_k_num+1]
    C_2 = (X_m_k_plus - X_m_k_minus) / (4 * X_m_k - 2 * X_m_k_minus -  2* X_m_k_plus)
    freq_estim_2 = C_2 * freq_bin + m_k ### No windowing

#     m_k_num = 3199
    if m_k_num == DC_num:
        print('DC Problem')
        phase_estim_2 = 0
#         plt.figure(1)
#         plt.plot(sig_magnitude)
# #         plt.plot(fit_x, fit_y)
#         plt.show()
#     phase_estim_2 = sig_phase[m_k_num]
#     freq_for_phase = freqline[m_k_num]
#     print(m_k_num)
    if 0:
        if len(sig) == 1280:
            m_k_num = 1900
#         elif len(sig) == 1024:
#             m_k_num = 4435
#         elif len(sig) == 3280:
#             m_k_num = 1159
    phase_estim_2 = sig_phase[m_k_num]
    freq_for_phase = freqline[m_k_num]

    
    return freq_estim_2, phase_estim_2, freqline, sig_magnitude, sig_phase, m_k_num, X_m_k, freq_for_phase

def FFT_interpolation_fix(sig, tau0, zero_num, DC_num, m_k_num_1280, m_k_num_1024): ###zero_num=1e5, DC_num=3000
    zero_padd = np.zeros(int(zero_num))
    sig_padd = np.concatenate((sig, zero_padd))
    N_data = len(sig)
    freqline, sig_FFT, sig_magnitude, sig_phase = FFT_cal(sig_padd, tau0)
    freq_bin = 1/ tau0 / len(sig_padd)
#     print(freq_bin)
    m_k_num = sig_magnitude[DC_num:80000].argmax()+DC_num ### frequency bin 序数 (3000 for 1e5)
#     print(m_k_num)
    m_k = freqline[m_k_num] ### 换算为实际频率
    X_m_k = sig_magnitude[m_k_num] ### magnitude (complex raw data)
    X_m_k_minus = sig_magnitude[m_k_num-1]
    X_m_k_plus = sig_magnitude[m_k_num+1]
    C_2 = (X_m_k_plus - X_m_k_minus) / (4 * X_m_k - 2 * X_m_k_minus -  2* X_m_k_plus)
    freq_estim_2 = C_2 * freq_bin + m_k ### No windowing
    
#     m_k_num = 3199
    if m_k_num == DC_num:
        print('DC Problem')
        freq_estim_2 = 0
#     phase_estim_2 = sig_phase[m_k_num]
#     freq_for_phase = freqline[m_k_num]
#     print(m_k_num)
    if 1:
        if len(sig) == 1280:
            m_k_num = m_k_num_1280
        elif len(sig) == 1024:
            m_k_num = m_k_num_1024
#         elif len(sig) == 3280:
#             m_k_num = 1159
#     m_k_num = m_k_num_1280
    
    phase_estim_2 = sig_phase[m_k_num]
    freq_for_phase = freqline[m_k_num]
    
    return freq_estim_2, phase_estim_2, freqline, sig_magnitude, sig_phase, m_k_num, X_m_k, freq_for_phase

def fit_func(x, a, b, c):
    return a*(x-b)**2 + c

def line_cal(line_data,pix_size=5.3e-6):
    DC_num = 400
#     pix_size = 5.3e-6
    freq_estim, phase_estim, freqline, sig_magnitude, sig_phase,  m_k_num, X_m_k, freq_for_phase = FFT_interpolation_2(line_data, pix_size, 1e5, DC_num)
    '''
        Fine Calc.
    '''
    FFT_start = np.where(sig_magnitude[DC_num:80000] > X_m_k*0.4)[0][0]+DC_num
    FFT_end = np.where(sig_magnitude[DC_num:80000] > X_m_k*0.4)[0][-1]+DC_num
    fit_x = freqline[FFT_start:FFT_end+1]
    fit_y = sig_magnitude[FFT_start:FFT_end+1]
#     plt.figure(1)
#     plt.plot(freqline,sig_magnitude)
# #     plt.plot(fit_x, fit_y)
# #     plt.plot(fit_x, gaus(fit_x, *popt))
#     plt.show()

#     params = curve_fit(fit_func, fit_x, fit_y, p0=[-1,5000,1], maxfev=80000)
#     [a, b, c] = params[0]
#     f_fit = b
    
    mean = sum(fit_x * fit_y) / sum(fit_y)
    sigma = np.sqrt(sum(fit_y * (fit_x - mean)**2) / sum(fit_y))
    popt,pcov = curve_fit(gaus,fit_x,fit_y,p0=[1,mean,sigma])
    f_fit = popt[1]
#     print(m_k_num, f_fit)
    

    '''
        Coarse Calc.
    '''
#     m_k = freqline[m_k_num] ### 换算为实际频率
#     X_m_k = sig_magnitude[m_k_num] ### magnitude (complex raw data)
#     X_m_k_minus = sig_magnitude[m_k_num-1]
#     X_m_k_plus = sig_magnitude[m_k_num+1]
#     C_2 = (X_m_k_plus - X_m_k_minus) / (4 * X_m_k - 2 * X_m_k_minus -  2* X_m_k_plus)
#     freq_estim_2 = C_2 * (freqline[1]-freqline[0]) + m_k ### No windowing
#     f_fit = freq_estim_2
#    plt.figure(1)
#    plt.plot(sig_magnitude)
##     plt.plot(fit_x, fit_y)
 #   plt.show()
    
    return f_fit, phase_estim, m_k_num

def line_cal_fix(line_data, m_k_num_1280, m_k_num_1024=0):
    DC_num = 400
    pix_size = 5.3e-6
    freq_estim, phase_estim, freqline, sig_magnitude, sig_phase,  m_k_num, X_m_k, freq_for_phase = FFT_interpolation_fix(line_data, pix_size, 1e5, DC_num, m_k_num_1280, m_k_num_1024)

    '''
        Fine Calc.
    '''
    FFT_start = np.where(sig_magnitude[DC_num:] > X_m_k*0.4)[0][0]+DC_num
    FFT_end = np.where(sig_magnitude[DC_num:] > X_m_k*0.4)[0][-1]+DC_num
    fit_x = freqline[FFT_start:FFT_end+1]
    fit_y = sig_magnitude[FFT_start:FFT_end+1]

#     params = curve_fit(fit_func, fit_x, fit_y, p0=[-1,5000,1], maxfev=80000)
#     print(params)
#     [a, b, c] = params[0]
# #     print(a,b,c)
#     f_fit = b
    mean = sum(fit_x * fit_y) / sum(fit_y)
    sigma = np.sqrt(sum(fit_y * (fit_x - mean)**2) / sum(fit_y))
    popt,pcov = curve_fit(gaus,fit_x,fit_y,p0=[1,mean,sigma])
    f_fit = popt[1]
    
    '''
        Coarse Calc.
    '''
#     m_k = freqline[m_k_num] ### 换算为实际频率
#     X_m_k = sig_magnitude[m_k_num] ### magnitude (complex raw data)
#     X_m_k_minus = sig_magnitude[m_k_num-1]
#     X_m_k_plus = sig_magnitude[m_k_num+1]
#     C_2 = (X_m_k_plus - X_m_k_minus) / (4 * X_m_k - 2 * X_m_k_minus -  2* X_m_k_plus)
#     freq_estim_2 = C_2 * (freqline[1]-freqline[0]) + m_k ### No windowing
#     f_fit = freq_estim_2
    
#     plt.figure(1)
#     plt.subplot(1,2,1)
#     plt.plot(line_data)
#     plt.subplot(1,2,2)
# #     print(m_k_num)
# #     plt.plot(sig_magnitude)
# #     plt.xlim(0,6000)
# #     plt.ylim(0,0.2)
#     plt.plot(fit_x, fit_y)
#     plt.plot(fit_x, a*(fit_x-b)**2 + c)
#     plt.show()
    
    return f_fit, phase_estim, m_k_num
def line_cal_test(line_data,pix_size=5.3e-6):
    DC_num = 0
#     pix_size = 5.3e-6
    freq_estim, phase_estim, freqline, sig_magnitude, sig_phase,  m_k_num, X_m_k, freq_for_phase = FFT_interpolation_2(line_data, pix_size, 1e5, DC_num)
    '''
        Fine Calc.
    '''
    FFT_start = np.where(sig_magnitude[DC_num:] > X_m_k*0.8)[0][0]+DC_num
    FFT_end = np.where(sig_magnitude[DC_num:] > X_m_k*0.8)[0][-1]+DC_num
    fit_x = freqline[FFT_start:FFT_end+1]
    fit_y = sig_magnitude[FFT_start:FFT_end+1]
#     plt.figure(1)
#     plt.plot(sig_magnitude)
#     plt.plot(fit_x, fit_y)
# #     plt.plot(fit_x, gaus(fit_x, *popt))
#     plt.show()

#     params = curve_fit(fit_func, fit_x, fit_y, p0=[-1,5000,1], maxfev=80000)
#     [a, b, c] = params[0]
#     f_fit = b
    
    mean = sum(fit_x * fit_y) / sum(fit_y)
    sigma = np.sqrt(sum(fit_y * (fit_x - mean)**2) / sum(fit_y))
    popt,pcov = curve_fit(gaus,fit_x,fit_y,p0=[1,mean,sigma])
    f_fit = popt[1]
#     print(m_k_num, f_fit)
    

    '''
        Coarse Calc.
    '''
#     m_k = freqline[m_k_num] ### 换算为实际频率
#     X_m_k = sig_magnitude[m_k_num] ### magnitude (complex raw data)
#     X_m_k_minus = sig_magnitude[m_k_num-1]
#     X_m_k_plus = sig_magnitude[m_k_num+1]
#     C_2 = (X_m_k_plus - X_m_k_minus) / (4 * X_m_k - 2 * X_m_k_minus -  2* X_m_k_plus)
#     freq_estim_2 = C_2 * (freqline[1]-freqline[0]) + m_k ### No windowing
#     f_fit = freq_estim_2
#    plt.figure(1)
#    plt.plot(sig_magnitude)
##     plt.plot(fit_x, fit_y)
 #   plt.show()
    
    return f_fit, phase_estim, m_k_num

def FFT_cal(data, tau0):
    N_data = len(data)
    # 频率轴设定,由实际点数与时间间隔决定，后半部分为负半轴
    freqline = np.fft.fftfreq(N_data, d=tau0)
    freqline = freqline[:(N_data+1)//2]
    sig_FFT = np.fft.fft(data) #two sides FFT
    sig_magnitude = (np.abs(sig_FFT[:(N_data+1)//2])) * 2 / N_data
    sig_phase = np.angle(sig_FFT)[:(N_data+1)//2]
    return freqline, sig_FFT, sig_magnitude, sig_phase

def FFT_interpolation_boxcar(data,tau0):
    N_data = len(data)
    freqline, sig_FFT, sig_magnitude, sig_phase = FFT_cal(data, tau0)
    freq_bin = 1/ tau0 / N_data
    m_k_num = sig_magnitude[:-1].argmax()+0 ### frequency bin 序数
    m_k = freqline[m_k_num] ### 换算为实际频率
    X_m_k = sig_FFT[m_k_num] ### magnitude (complex raw data)
    X_m_k_minus = sig_FFT[m_k_num-1]
    X_m_k_plus = sig_FFT[m_k_num+1]
    C_1 = (X_m_k_minus - X_m_k_plus) / (2 * X_m_k - X_m_k_minus - X_m_k_plus)
    freq_estim_boxcar = np.real(C_1) * freq_bin + m_k ### No windowing
    
    X_m_k = sig_magnitude[m_k_num] ### magnitude (complex raw data)
    X_m_k_minus = sig_magnitude[m_k_num-1]
    X_m_k_plus = sig_magnitude[m_k_num+1]
    C_2 = (X_m_k_plus - X_m_k_minus) / (4 * X_m_k - 2 * X_m_k_minus -  2* X_m_k_plus)
    freq_estim_2 = C_2 * freq_bin + m_k ### No windowing
    
    Period_num = freq_estim_2 * N_data * tau0
    N_sel = int(N_data * int(Period_num) / Period_num)
    data = data[:N_sel+1]
    freqline_est, sig_FFT_est, sig_freq_est, sig_phase_est = FFT_cal(data, tau0)
    phase_estim_boxcar = sig_phase_est[sig_freq_est.argmax()]
#     print(freqline[sig_freq_est.argmax()]) ### Real Freq used in Phase calculation
    
    return freq_estim_boxcar, phase_estim_boxcar, freqline_est, sig_freq_est, freqline_est[sig_freq_est.argmax()]

def FFT_interpolation_nonlinearity_compare(data,tau0):
    N_data = len(data)
    freqline, sig_FFT, sig_magnitude, sig_phase = FFT_cal(data, tau0)
    freq_bin = 1/ tau0 / N_data
    m_k_num = sig_magnitude[20:-1].argmax()+20 ### frequency bin 序数
    m_k = freqline[m_k_num] ### 换算为实际频率
    X_m_k = sig_FFT[m_k_num] ### magnitude (complex raw data)
    X_m_k_minus = sig_FFT[m_k_num-1]
    X_m_k_plus = sig_FFT[m_k_num+1]
    C_1 = (X_m_k_minus - X_m_k_plus) / (2 * X_m_k - X_m_k_minus - X_m_k_plus)
    freq_estim_boxcar = np.real(C_1) * freq_bin + m_k ### No windowing
    
    Period_num = freq_estim_boxcar * N_data * tau0
    N_sel = int(N_data * int(Period_num) / Period_num)
    data = data[:N_sel+1]
    freqline_est, sig_FFT_est, sig_freq_est, sig_phase_est = FFT_cal(data, tau0)
#     phase_estim_boxcar = sig_phase_est[sig_freq_est.argmax()]
#     print(freqline[sig_freq_est.argmax()]) ### Real Freq used in Phase calculation
    return freq_estim_boxcar, sig_phase_est, freqline_est, sig_freq_est


def FFT_interpolation_compare(data,tau0):
    N_data = len(data)
    freqline, sig_FFT, sig_magnitude, sig_phase = FFT_cal(data, tau0)
    freq_bin = 1/ tau0 / N_data
    
    m_k_num = sig_magnitude[3000:].argmax()+3000 ### frequency bin 序数 (3000 for 1e5)
    print(m_k_num)
    m_k = freqline[m_k_num] ### 换算为实际频率
    X_m_k = sig_FFT[m_k_num] ### magnitude (complex raw data)
    X_m_k_minus = sig_FFT[m_k_num-1]
    X_m_k_plus = sig_FFT[m_k_num+1]
    C_1 = (X_m_k_minus - X_m_k_plus) / (2 * X_m_k - X_m_k_minus - X_m_k_plus)
    freq_estim_1 = np.real(C_1) * freq_bin + m_k ### No windowing
    
    X_m_k = sig_magnitude[m_k_num] ### magnitude (complex raw data)
    X_m_k_minus = sig_magnitude[m_k_num-1]
    X_m_k_plus = sig_magnitude[m_k_num+1]
    C_2 = (X_m_k_plus - X_m_k_minus) / (4 * X_m_k - 2 * X_m_k_minus -  2* X_m_k_plus)
    freq_estim_2 = C_2 * freq_bin + m_k ### No windowing
    
    X_m_k = sig_magnitude[m_k_num] ### magnitude (complex raw data)
    X_m_k_minus = sig_magnitude[m_k_num-1]
    X_m_k_plus = sig_magnitude[m_k_num+1]
    C_3 = P*(X_m_k_plus - X_m_k_minus) / (X_m_k + X_m_k_minus + X_m_k_plus)
    freq_estim_3 = C_3 * freq_bin + m_k ### No windowing
    
    X_m_k = sig_FFT[m_k_num] ### magnitude (complex raw data)
    X_m_k_minus = sig_FFT[m_k_num-1]
    X_m_k_plus = sig_FFT[m_k_num+1]
    C_4 = Q*(X_m_k_minus - X_m_k_plus) / (2 * X_m_k + X_m_k_minus + X_m_k_plus)
    freq_estim_4 = np.real(C_4) * freq_bin + m_k ### No windowing
    
    return freq_estim_1, freq_estim_2, freq_estim_3, freq_estim_4

def FFT_interpolation_1(sig, tau0, zero_num, DC_num): ###zero_num=1e5, DC_num=3000
    zero_padd = np.zeros(int(zero_num))
    sig_padd = np.concatenate((zero_padd, sig, zero_padd))
    N_data = len(sig)
    freqline, sig_FFT, sig_magnitude, sig_phase = FFT_cal(sig_padd, tau0)
    freq_bin = 1/ tau0 / len(sig_padd)
    m_k_num = sig_magnitude[DC_num:].argmax()+DC_num ### frequency bin 序数 (3000 for 1e5)
    m_k = freqline[m_k_num] ### 换算为实际频率
    X_m_k = sig_FFT[m_k_num] ### magnitude (complex raw data)
    X_m_k_minus = sig_FFT[m_k_num-1]
    X_m_k_plus = sig_FFT[m_k_num+1]
    C_1 = (X_m_k_minus - X_m_k_plus) / (2 * X_m_k - X_m_k_minus - X_m_k_plus)
    freq_estim_1 = np.real(C_1) * freq_bin + m_k ### No windowing
    Period_num = freq_estim_1 * len(sig) * tau0
    N_sel = int(len(sig) * int(Period_num) / Period_num)
    sig_sel = sig[:N_sel+1]
    sig_sel_padd = np.concatenate((zero_padd, sig_sel, zero_padd))
    freqline_est, sig_FFT_est, sig_magnitude_est, sig_phase_est = FFT_cal(sig_sel_padd, tau0)
    phase_estim_1 = sig_phase_est[sig_magnitude_est[DC_num:].argmax()+DC_num]
    if phase_estim_1 < 0:
        phase_estim_1 += np.pi
    return freq_estim_1, phase_estim_1


def FFT_interpolation_3(sig, tau0, zero_num, DC_num): ###zero_num=1e5, DC_num=3000
    zero_padd = np.zeros(int(zero_num))
    sig_padd = np.concatenate((zero_padd, sig, zero_padd))
    N_data = len(sig)
    freqline, sig_FFT, sig_magnitude, sig_phase = FFT_cal(sig_padd, tau0)
    freq_bin = 1/ tau0 / len(sig_padd)
    m_k_num = sig_magnitude[DC_num:].argmax()+DC_num ### frequency bin 序数 (3000 for 1e5)
    m_k = freqline[m_k_num] ### 换算为实际频率
    X_m_k = sig_magnitude[m_k_num] ### magnitude (complex raw data)
    X_m_k_minus = sig_magnitude[m_k_num-1]
    X_m_k_plus = sig_magnitude[m_k_num+1]
    C_3 = P*(X_m_k_plus - X_m_k_minus) / (X_m_k + X_m_k_minus + X_m_k_plus)
    freq_estim_3 = C_3 * freq_bin + m_k ### No windowing
    Period_num = freq_estim_3 * len(sig) * tau0
    N_sel = int(len(sig) * int(Period_num) / Period_num)
    sig_sel = sig[:N_sel+1]
    sig_sel_padd = np.concatenate((zero_padd, sig_sel, zero_padd))
    freqline_est, sig_FFT_est, sig_magnitude_est, sig_phase_est = FFT_cal(sig_sel_padd, tau0)
    phase_estim_3 = sig_phase_est[sig_magnitude_est[DC_num:].argmax()+DC_num]
    if phase_estim_3 < 0:
        phase_estim_3 += np.pi
    return freq_estim_3, phase_estim_3

def FFT_interpolation_4(sig, tau0, zero_num, DC_num): ###zero_num=1e5, DC_num=3000
    zero_padd = np.zeros(int(zero_num))
    sig_padd = np.concatenate((zero_padd, sig, zero_padd))
    N_data = len(sig)
    freqline, sig_FFT, sig_magnitude, sig_phase = FFT_cal(sig_padd, tau0)
    freq_bin = 1/ tau0 / len(sig_padd)
    m_k_num = sig_magnitude[DC_num:].argmax()+DC_num ### frequency bin 序数 (3000 for 1e5)
    m_k = freqline[m_k_num] ### 换算为实际频率
    X_m_k = sig_FFT[m_k_num] ### magnitude (complex raw data)
    X_m_k_minus = sig_FFT[m_k_num-1]
    X_m_k_plus = sig_FFT[m_k_num+1]
    C_4 = Q*(X_m_k_minus - X_m_k_plus) / (2 * X_m_k + X_m_k_minus + X_m_k_plus)
    freq_estim_4 = np.real(C_4) * freq_bin + m_k ### No windowing
    Period_num = freq_estim_4 * len(sig) * tau0
    N_sel = int(len(sig) * int(Period_num) / Period_num)
    sig_sel = sig[:N_sel+1]
    sig_sel_padd = np.concatenate((zero_padd, sig_sel, zero_padd))
    freqline_est, sig_FFT_est, sig_magnitude_est, sig_phase_est = FFT_cal(sig_sel_padd, tau0)
    phase_estim_4 = sig_phase_est[sig_magnitude_est[DC_num:].argmax()+DC_num]
    if phase_estim_4 < 0:
        phase_estim_4 += np.pi
    return freq_estim_4, phase_estim_4

# tau0 = 0.001
# fs = 1/tau0
# x = np.arange(0+tau0, 1, tau0)



